#!/usr/bin/python3

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches

# date, priceOpen, priceHigh, priceLow, priceClose, volume, volumePrice
sample_data_lines = " \n \
2011/10/10	6.25	6.29	6.14	6.15	32428434	274365568.00 \n \
2011/10/11	6.51	6.51	6.21	6.30	59130344	511050464.00 \n \
2011/10/12	6.27	6.65	6.22	6.58	89412222	787072832.00 \n \
2011/10/13	6.55	6.67	6.49	6.60	51298778	453963296.00 \n \
2011/10/14	6.57	6.61	6.46	6.56	40504198	356399904.00 \n \
2011/10/17	6.57	6.78	6.55	6.65	46009758	410848352.00 \n \
2011/10/18	6.58	6.63	6.48	6.51	41115344	362177728.00 \n \
2011/10/19	6.51	6.65	6.48	6.57	36188377	319545728.00 \n \
2011/10/20	6.50	6.59	6.37	6.46	50658455	443175232.00 \n \
2011/10/21	6.48	6.70	6.43	6.62	59306969	526976704.00 \n \
2011/10/24	6.64	6.94	6.60	6.89	94932084	864185472.00 \n \
2011/10/25	6.91	6.98	6.81	6.89	66314936	608223040.00 \n \
2011/10/26	6.86	7.03	6.78	6.91	68454846	631301568.00 \n \
2011/10/27	6.99	7.08	6.94	6.99	45585827	422360448.00 \n \
"

def get_data_lines_from_file(fpath):
    with open(fpath, 'r') as f:
        return '\n'.join(f.readlines())

def get_prices_volumes(data_lines):
    """
    Return (prices, volumes)
    prices: a list of [pO, pH, pL, pC]
    volumes: a list of volume
    """
    prices = []
    volumes = []
    for line in data_lines.split('\n'):
        if line.strip() == "":
            continue
        date, pO, pH, pL, pC, v, vP = line.split()
        prices.append([pO, pH, pL, pC])
        volumes.append(vP)
    return (prices, volumes)

def test_draw_klines(fig, prices, withcolor=True):
    """
    Draw out k line according to prices.
    prices: [[pO, pH, pL, pC], ...]
    """

    x_limit_low = 0
    x_limit_high = len(prices)
    y_limit_low = -1
    y_limit_high = -1

    width = 1
    
    ax1 = fig.add_subplot(211)

    for idx, price in enumerate(prices):
        pO, pH, pL, pC = [float(p) for p in price]

        p_up = pO if pO > pC else pC
        p_down = pO if pO <= pC else pC
        start_point_rec = (idx, pO if pO < pC else pC)
        height_rec = p_up - p_down
        start_point_line_down = (idx + 0.5 * width, pL)
        height_down = p_down - pL
        start_point_line_up = (idx + 0.5 * width, p_up)
        height_up = pH - p_up

        if withcolor:
            color = 'red' if (pC > pO) else 'green'
        else:
            color = 'none'
        ax1.add_patch(patches.Rectangle(start_point_rec, 1, height_rec, alpha=1, facecolor=color))
        ax1.add_patch(patches.Rectangle(start_point_line_down, 0, height_down, alpha=1, facecolor='none'))
        ax1.add_patch(patches.Rectangle(start_point_line_up, 0, height_up, alpha=1, facecolor='none'))

        # deal with y_limit_low and y_limit_high
        if pL < y_limit_low or y_limit_low == -1:
            y_limit_low = pL
        if pH > y_limit_high or y_limit_high == -1:
            y_limit_high = pH

    plt.xlim(int(x_limit_low), int(x_limit_high) + 1)
    plt.ylim(y_limit_low * 0.9, y_limit_high * 1.1)

def test_draw_volumes(fig, volumes, prices, withcolor=True):
    """
    Draw out volumes bars accoring volume prices.
    """
    ax1 = fig.add_subplot(212)

    x_limit_low = 0
    x_limit_high = len(volumes)
    y_limit_low = 0
    y_limit_high = -1 # to be determined

    for idx, volume in enumerate(volumes):
        volume = float(volume)
        pO, pH, pL, pC = [float(p) for p in prices[idx]]
        if withcolor:
            color = 'red' if (pC > pO) else 'green'
        else:
            color = 'none'
        start_point = (idx, 0)
        ax1.add_patch(patches.Rectangle(start_point, 1, volume, alpha=1, facecolor=color))
        # deal with  y_limit_high
        if volume > y_limit_high or y_limit_high == -1:
            y_limit_high = volume

    plt.xlim(int(x_limit_low), int(x_limit_high) + 1)
    plt.ylim(y_limit_low, y_limit_high * 1.1)


if __name__ == "__main__":
    import sys

    data_type = "file"

    # data type
    if len(sys.argv) > 1:
        data_type = sys.argv[1]
    # figure size
    #fig_size = (10, 10)
    fig_size = None

    if data_type == "simple":
        prices, volumes = get_prices_volumes(sample_data_lines)
    else:
        data_lines = get_data_lines_from_file('SH#600000.txt')
        prices, volumes = get_prices_volumes(data_lines)

    fig = plt.figure(figsize=fig_size)
    test_draw_klines(fig, prices, False)
    test_draw_volumes(fig, volumes, prices, False)
    plt.show()
